/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.sharding.client.interceptor;

import com.lvyh.lightframe.sharding.client.common.DataTableInfo;
import com.lvyh.lightframe.sharding.client.common.ShardModel;
import com.lvyh.lightframe.sharding.client.context.ShardingContext;
import com.lvyh.lightframe.sharding.client.parser.SimpleSqlParser;
import com.lvyh.lightframe.sharding.client.strategy.ShardingStrategyFactory;
import com.lvyh.lightframe.sharding.client.util.ReflectionUtil;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.ReflectorFactory;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.factory.ObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.reflection.wrapper.ObjectWrapperFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.Objects;
import java.util.Properties;

@Component
@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
public class TableUpdateInterceptor implements Interceptor {
    private Logger logger = LoggerFactory.getLogger(this.getClass());
    private static final ObjectFactory DEFAULT_OBJECT_FACTORY = new DefaultObjectFactory();
    private static final ObjectWrapperFactory DEFAULT_OBJECT_WRAPPER_FACTORY = new DefaultObjectWrapperFactory();
    private static final ReflectorFactory REFLECTOR_FACTORY = new DefaultReflectorFactory();

    @Autowired
    private ShardingStrategyFactory shardingStrategyFactory;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        StatementHandler statement = (StatementHandler) invocation.getTarget();
        Field delegateField = ReflectionUtil.getField(statement, "delegate");
        StatementHandler prepared = (StatementHandler) delegateField.get(statement);
        Field msField = ReflectionUtil.getField(prepared, "mappedStatement");
        MappedStatement ms = (MappedStatement) msField.get(prepared);
        String sqlCommandType = ms.getSqlCommandType().name();
        if (Objects.equals(SqlCommandType.SELECT.name(), sqlCommandType)) {
            return invocation.proceed();
        }

        // Get the ID of the node, that is, the ID of the SQL statement
        String sqlId = ms.getId();
        BoundSql boundSql = prepared.getBoundSql();
        logger.info("[TableUpdateInterceptor] print mapped statement id: {}", sqlId);

        Object parameterObject = boundSql.getParameterObject();
        ShardModel shardModel = ShardingContext.get();

        if (shardModel != null) {
            this.rewrite(ms, boundSql, shardModel, parameterObject);
        }
        return invocation.proceed();
    }


    private void rewrite(MappedStatement ms, BoundSql boundSql, ShardModel shardModel, Object parameter) throws Exception {
        String logicDsKey = shardModel.getLogicDatasourceName();
        String logicTableName = shardModel.getLogicTableName();
        String shardingColumn = shardModel.getShardKey();
        String shardingAlg = shardModel.getShardAlg();
        int tableShardingNum = shardModel.getTableShardNum();
        int dbShardingNum = shardModel.getDbShardNum();

        String sql = boundSql.getSql();
        String sqlId = ms.getId();
        Object shardingValue = null;
        if (sqlId.endsWith("selectByExample_COUNT") || sqlId.endsWith("selectByExample")) {
            shardingValue = computeShardingValue(shardingColumn, parameter);
        } else
            shardingValue = ReflectionUtil.getFieldValue(parameter, shardingColumn);

        if (StringUtils.isEmpty(shardingValue)) {
            return;
        }

        DataTableInfo dataTableInfo = this.shardingStrategyFactory.getShardInfo(shardingAlg, logicDsKey, logicTableName, shardingValue, tableShardingNum, dbShardingNum);
        String newSql = SimpleSqlParser.getConvertedSql(logicTableName, dataTableInfo.getPhysicalTableName(), sql);
        ReflectionUtil.setField(boundSql, "sql", newSql);

        logger.info("[TableUpdateInterceptor] print current data table：{}", dataTableInfo.toString());
    }

    private MetaObject getMetaObject(Object parameterObject) {
        MetaObject metaObject = MetaObject.forObject(parameterObject, DEFAULT_OBJECT_FACTORY, DEFAULT_OBJECT_WRAPPER_FACTORY, REFLECTOR_FACTORY);
        return metaObject;
    }

    private String computeShardingValue(String shardingColumn, Object parameterObject) {
        // Reflection get parameter object
        MetaObject param = getMetaObject(parameterObject);
        Object oredCriteriaList = param.getValue("oredCriteria");
        if (oredCriteriaList != null && oredCriteriaList instanceof ArrayList) {
            Object criteria = ((ArrayList) oredCriteriaList).get(0);
            MetaObject criteriaBaseObject = getMetaObject(criteria);
            Object criteriaList = criteriaBaseObject.getValue("criteria");
            if (criteriaList != null && criteriaList instanceof ArrayList) {
                Object criteriaObject = ((ArrayList) criteriaList).get(0);
                MetaObject criteriaMetaObject = getMetaObject(criteriaObject);
                Object condition = criteriaMetaObject.getValue("condition");
                Object value = criteriaMetaObject.getValue("value");
                if (condition instanceof String && condition.toString().contains(shardingColumn)) {
                    return value.toString();
                }
            }
        }
        return "";
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }
}
